package org.dxy.trigger;

import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.api.ndarray.INDArray;

/**
 * 简单的ND4J测试 - 自动检测GPU/CPU后端
 */
public class SimpleGPUTest {
    public static void main(String[] args) {
        System.out.println("=== ND4J Backend Detection Test ===");
        
        // 不设置任何属性，让ND4J自动选择最佳后端
        try {
            // 获取默认后端信息
            String backendName = Nd4j.getBackend().getClass().getSimpleName();
            System.out.println("Default Backend: " + backendName);
            
            boolean isGPU = backendName.toLowerCase().contains("cuda") || 
                           backendName.toLowerCase().contains("gpu");
            System.out.println("GPU Mode: " + (isGPU ? "✅ YES" : "❌ NO (CPU)"));
            
            // 获取设备数量
            try {
                int devices = Nd4j.getAffinityManager().getNumberOfDevices();
                System.out.println("Available devices: " + devices);
            } catch (Exception e) {
                System.out.println("Device count: Unknown");
            }
            
            // 性能测试
            performanceTest(isGPU ? "GPU" : "CPU", isGPU ? 1000 : 500);
            
        } catch (Exception e) {
            System.err.println("❌ Backend test failed: " + e.getMessage());
            e.printStackTrace();
        }
    }
    
    private static void performanceTest(String mode, int size) {
        try {
            System.out.println("\n=== " + mode + " Performance Test ===");
            System.out.println("Creating " + size + "x" + size + " matrices...");
            
            long startTime = System.currentTimeMillis();
            
            INDArray a = Nd4j.rand(size, size);
            INDArray b = Nd4j.rand(size, size);
            INDArray result = a.mmul(b);
            
            long endTime = System.currentTimeMillis();
            
            System.out.println("Matrix multiplication completed in: " + (endTime - startTime) + "ms");
            System.out.println("Result shape: " + java.util.Arrays.toString(result.shape()));
            System.out.println("First element: " + String.format("%.6f", result.getDouble(0, 0)));
            
            // 检查结果设备
            try {
//                System.out.println("Result device ID: " + result.data().getDeviceId());
            } catch (Exception e) {
                System.out.println("Device ID: Not available");
            }
            
            System.out.println("✅ " + mode + " test completed successfully!");
            
        } catch (Exception e) {
            System.err.println("❌ " + mode + " performance test failed: " + e.getMessage());
        }
    }
}